package com.kehao.leetcode.hard.tree_graph;


import org.junit.Test;

import java.util.HashMap;
import java.util.HashSet;

public class LowestCommonAncestorSolution {
    private static class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;
        TreeNode(int x) { val = x; }
    }
    private TreeNode createTree(Integer[] array){
        TreeNode[] treeNodes = new TreeNode[array.length];
        for (int i = 0; i < array.length; i++) {
            if(array[i]!=null){
                treeNodes[i] = new TreeNode(array[i]);
            }
        }
        for (int i = 0; i < array.length / 2 ; i++) {
            if(2*i+1<array.length)treeNodes[i].left = treeNodes[2*i+1];
            if(2*i+2<array.length)treeNodes[i].right = treeNodes[2*i+2];
        }
        return treeNodes[0];
    }
    private TreeNode getTreeNodeByNum(TreeNode root,int num){
        if(root == null||root.val==num) return root;
        TreeNode left = getTreeNodeByNum(root.left, num);
        return left !=null? left:getTreeNodeByNum(root.right, num);
    }

    public TreeNode lowestCommonAncestor(TreeNode root, TreeNode p, TreeNode q) {
        HashMap<TreeNode,TreeNode> sonParentMap = new HashMap<>();
        fillMap(root,sonParentMap);
        HashSet<TreeNode> vistedSet = new HashSet<>();
        while (sonParentMap.containsKey(p)&&sonParentMap.containsKey(q)){
            if(vistedSet.contains(p)) {
                return p;
            }else {
                vistedSet.add(p);
                p = sonParentMap.get(p);
            }
            if(vistedSet.contains(q)) {
                return q;
            }else {
                vistedSet.add(q);
                q = sonParentMap.get(q);
            }
        }
        while (sonParentMap.containsKey(p)){
            if(vistedSet.contains(p)) {
                return p;
            }else {
                vistedSet.add(p);
                p = sonParentMap.get(p);
            }
        }
        while (sonParentMap.containsKey(q)){
            if(vistedSet.contains(q)) {
                return q;
            }else {
                vistedSet.add(q);
                q = sonParentMap.get(q);
            }
        }
        return root;
    }

    private void fillMap(TreeNode root, HashMap<TreeNode, TreeNode> sonParentMap) {
        if(root.left!=null) {
            sonParentMap.put(root.left,root);
            fillMap(root.left,sonParentMap);
        }
        if(root.right!=null) {
            sonParentMap.put(root.right,root);
            fillMap(root.right,sonParentMap);
        }
    }

    @Test
    public void test01(){
        Integer[] numArray = {3,5,1,6,2,0,8,null,null,7,4};
        TreeNode root = createTree(numArray);
        TreeNode p = getTreeNodeByNum(root,5);
        TreeNode q = getTreeNodeByNum(root,1);
        System.out.println(lowestCommonAncestor(root,p,q).val);
    }

    @Test
    public void test02(){
        Integer[] numArray = {3,5,1,6,2,0,8,null,null,7,4};
        TreeNode root = createTree(numArray);
        TreeNode p = getTreeNodeByNum(root,5);
        TreeNode q = getTreeNodeByNum(root,4);
        System.out.println(lowestCommonAncestor(root,p,q).val);
    }

    @Test
    public void test03(){
        Integer[] numArray = {1,2};
        TreeNode root = createTree(numArray);
        TreeNode p = getTreeNodeByNum(root,1);
        TreeNode q = getTreeNodeByNum(root,2);
        System.out.println(lowestCommonAncestor(root,p,q).val);
    }
}
